#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
机器人前向动力学程序
基于参考代码简化，仅包含前向动力学和滑块控制功能
"""

import numpy as np
import os
import sys
import time
import tkinter as tk
from tkinter import ttk, Frame
import threading
import pybullet as p
import pybullet_data

# 全局变量
robot_id = None  # 机器人ID
joint_ids = []   # 关节ID列表
joint_names = [] # 关节名称列表
running = True   # 控制仿真运行状态

# URDF文件路径 - 使用绝对路径
urdf_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "robot2.urdf")

class RobotController:
    """机器人控制器类，管理机器人模型和前向动力学"""
    
    def __init__(self, urdf_path):
        """初始化机器人控制器
        
        参数:
            urdf_path: URDF文件路径
        """
        self.urdf_path = urdf_path
        self.robot_id = None
        self.joint_ids = []
        self.joint_names = []
        self.joint_angles = []
        
        # 检查URDF文件是否存在
        if not os.path.exists(urdf_path):
            raise FileNotFoundError(f"错误: URDF文件不存在: {urdf_path}")
    
    def init_pybullet(self):
        """初始化PyBullet仿真环境"""
        # 连接到PyBullet
        try:
            p.connect(p.GUI)
        except p.error:
            print("警告: 无法连接到GUI模式，回退到DIRECT模式")
            p.connect(p.DIRECT)
            
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.setGravity(0, 0, -9.8)
        
        # 加载地面
        plane_id = p.loadURDF('plane.urdf')
        p.changeDynamics(plane_id, -1, lateralFriction=1)
        
        # 加载机器人
        start_pos = [0, 0, 0.5]
        start_ori = p.getQuaternionFromEuler([0, 0, 0])
        self.robot_id = p.loadURDF(self.urdf_path, start_pos, start_ori, useFixedBase=True)
        
        # 获取关节信息
        num_joints = p.getNumJoints(self.robot_id)
        print(f"机器人关节数量: {num_joints}")
        
        # 收集所有可控制关节
        self.joint_ids = []
        self.joint_names = []
        
        for i in range(num_joints):
            info = p.getJointInfo(self.robot_id, i)
            joint_name = info[1].decode('utf-8')
            joint_type = info[2]
            
            # 只收集可移动的关节
            if joint_type != p.JOINT_FIXED:
                self.joint_ids.append(i)
                self.joint_names.append(joint_name)
                print(f"关节 {i}: {joint_name}, 类型: {joint_type}")
        
        # 初始化关节角度为0
        self.joint_angles = [0.0] * len(self.joint_ids)
        
        # 设置仿真步长
        p.setTimeStep(1/240)
        
        # 设置摄像机
        p.resetDebugVisualizerCamera(
            cameraDistance=1.0,
            cameraYaw=100,
            cameraPitch=-30,
            cameraTargetPosition=[0, 0, 0.5]
        )
        
        return self.robot_id, self.joint_ids, self.joint_names
    
    def apply_joint_angles(self):
        """应用关节角度到机器人"""
        for i, joint_id in enumerate(self.joint_ids):
            p.setJointMotorControl2(
                self.robot_id,
                joint_id,
                p.POSITION_CONTROL,
                targetPosition=self.joint_angles[i],
                force=500
            )
    
    def reset_joint_angles(self):
        """重置所有关节角度为0"""
        self.joint_angles = [0.0] * len(self.joint_ids)
        self.apply_joint_angles()


class SliderInterface:
    """滑块界面类，提供用户界面控制关节角度"""
    
    def __init__(self, robot_controller):
        """初始化滑块界面
        
        参数:
            robot_controller: RobotController实例
        """
        self.robot_controller = robot_controller
        
        # 创建主窗口
        self.root = tk.Tk()
        self.root.title("机器人前向动力学控制面板")
        self.root.geometry("800x600")
        
        # 创建滑块变量
        self.joint_vars = []
        for _ in range(len(robot_controller.joint_ids)):
            self.joint_vars.append(tk.DoubleVar(value=0.0))
        
        # 创建界面组件
        self.create_widgets()
        
        # 绑定关闭事件
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)
    
    def create_widgets(self):
        """创建界面组件"""
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建控制区域
        control_frame = ttk.LabelFrame(main_frame, text="关节控制", padding="5")
        control_frame.pack(fill=tk.BOTH, expand=True, pady=5)
        
        # 创建滑块
        for i, joint_name in enumerate(self.robot_controller.joint_names):
            # 创建行框架
            row_frame = ttk.Frame(control_frame)
            row_frame.pack(fill=tk.X, pady=2)
            
            # 关节名称标签
            ttk.Label(row_frame, text=f"{joint_name}:", width=15).pack(side=tk.LEFT)
            
            # 角度值标签
            value_label = ttk.Label(row_frame, text="0.00", width=8)
            value_label.pack(side=tk.LEFT)
            
            # 滑块
            scale = ttk.Scale(
                row_frame,
                from_=-3.14,  # -180度
                to=3.14,      # 180度
                orient=tk.HORIZONTAL,
                variable=self.joint_vars[i],
                length=400
            )
            scale.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)
            
            # 绑定值变化事件
            self.joint_vars[i].trace_add("write", lambda *args, idx=i, label=value_label: 
                                        self.on_slider_change(idx, label))
        
        # 创建按钮区域
        button_frame = ttk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=10)
        
        # 重置按钮
        ttk.Button(
            button_frame,
            text="重置所有关节",
            command=self.reset_joints
        ).pack(side=tk.RIGHT, padx=5)
    
    def on_slider_change(self, index, label):
        """滑块值变化处理
        
        参数:
            index: 关节索引
            label: 值标签
        """
        # 获取新值
        value = self.joint_vars[index].get()
        
        # 更新标签
        label.config(text=f"{value:.2f}")
        
        # 更新控制器中的关节角度
        self.robot_controller.joint_angles[index] = value
    
    def reset_joints(self):
        """重置所有关节"""
        # 重置控制器中的关节角度
        self.robot_controller.reset_joint_angles()
        
        # 重置所有滑块值
        for i, var in enumerate(self.joint_vars):
            var.set(0.0)
    
    def update_ui(self):
        """更新界面"""
        # 在这里可以添加界面更新逻辑
        pass
    
    def on_closing(self):
        """关闭窗口处理"""
        global running
        running = False
        self.root.destroy()
    
    def start(self):
        """启动界面"""
        # 启动主循环
        self.root.mainloop()


def simulation_thread(robot_controller, ui):
    """仿真线程
    
    参数:
        robot_controller: RobotController实例
        ui: SliderInterface实例
    """
    global running
    
    # 摄像机参数
    camera_distance = 1.0
    camera_yaw = 50
    camera_pitch = -30
    camera_target = [0, 0, 0.5]
    
    # 主循环
    while running:
        try:
            # 检查界面是否关闭
            if not hasattr(ui, 'root') or not ui.root.winfo_exists():
                break
            
            # 处理键盘事件 - 摄像头控制
            keys = p.getKeyboardEvents()
            if p.B3G_LEFT_ARROW in keys and (keys[p.B3G_LEFT_ARROW] & p.KEY_IS_DOWN):
                camera_yaw -= 2  # 左键减小yaw值，摄像头向左转
            if p.B3G_RIGHT_ARROW in keys and (keys[p.B3G_RIGHT_ARROW] & p.KEY_IS_DOWN):
                camera_yaw += 2  # 右键增加yaw值，摄像头向右转
            
            # 更新摄像机位置
            p.resetDebugVisualizerCamera(
                cameraDistance=camera_distance,
                cameraYaw=camera_yaw,
                cameraPitch=camera_pitch,
                cameraTargetPosition=camera_target
            )
            
            # 应用关节角度
            robot_controller.apply_joint_angles()
            
            # 步进仿真
            p.stepSimulation()
            time.sleep(1/240)
            
        except KeyboardInterrupt:
            print("仿真线程被中断")
            break
        except Exception as e:
            print(f"仿真线程出错: {e}")
            break
    
    # 断开PyBullet连接
    p.disconnect()


def main():
    """主函数"""
    global running
    
    try:
        # 创建机器人控制器
        robot_controller = RobotController(urdf_file)
        
        # 初始化PyBullet
        robot_id, joint_ids, joint_names = robot_controller.init_pybullet()
        
        # 创建滑块界面
        ui = SliderInterface(robot_controller)
        
        # 启动仿真线程
        sim_thread = threading.Thread(
            target=simulation_thread,
            args=(robot_controller, ui),
            daemon=True
        )
        sim_thread.start()
        
        # 启动界面
        ui.start()
        
    except KeyboardInterrupt:
        print("程序被中断")
    except Exception as e:
        print(f"程序出错: {e}")
    finally:
        # 确保断开PyBullet连接
        running = False
        if p.isConnected():
            p.disconnect()


if __name__ == "__main__":
    main()
